# coding=utf-8
"""
"""

import PIL
from PIL import ImageFont
from PIL import Image
from PIL import ImageDraw
import cv2;
import numpy as np;
import os;
from math import *
import sys
import matplotlib.pyplot as plt

index = {'0': 0,
         '1': 1,
         '2': 2,
         '3': 3,
         '4': 4,
         '5': 5,
         '6': 6,
         '7': 7,
         '8': 8,
         '9': 9,
         '-': 10}
chars = ["0", "1", "2", "3", "4", "5","6","7","8","9","-"]

def rot(img, angel, shape, max_angel):
    """
        添加放射畸变
        img 输入图像
        factor 畸变的参数
        size 为图片的目标尺寸
    """
    size_o = [shape[1], shape[0]]
    size = (shape[1] + int(shape[0] * cos((float(max_angel) / 180) * 3.14)), shape[0])
    interval = abs(int(sin((float(angel) / 180) * 3.14) * shape[0]));
    pts1 = np.float32([[0, 0], [0, size_o[1]], [size_o[0], 0], [size_o[0], size_o[1]]])
    if (angel > 0):
        pts2 = np.float32([[interval, 0], [0, size[1]], [size[0], 0], [size[0] - interval, size_o[1]]])
    else:
        pts2 = np.float32([[0, 0], [interval, size[1]], [size[0] - interval, 0], [size[0], size_o[1]]])
    M = cv2.getPerspectiveTransform(pts1, pts2);
    dst = cv2.warpPerspective(img, M, size);
    return dst


def rotRandrom(img, factor, size):
    """
    添加透视畸变
    """
    shape = size;
    pts1 = np.float32([[0, 0], [0, shape[0]], [shape[1], 0], [shape[1], shape[0]]])
    pts2 = np.float32([[r(factor), r(factor)], [r(factor), shape[0] - r(factor)], [shape[1] - r(factor), r(factor)],
                       [shape[1] - r(factor), shape[0] - r(factor)]])
    M = cv2.getPerspectiveTransform(pts1, pts2);
    dst = cv2.warpPerspective(img, M, size);
    return dst


def tfactor(img):
    """
    添加饱和度光照的噪声
    """
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV);
    hsv[:, :, 0] = hsv[:, :, 0] * (0.8 + np.random.random() * 0.2);
    hsv[:, :, 1] = hsv[:, :, 1] * (0.3 + np.random.random() * 0.7);
    hsv[:, :, 2] = hsv[:, :, 2] * (0.2 + np.random.random() * 0.8);

    img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR);
    return img


def random_envirment(img, data_set):
    """
    添加自然环境的噪声
    """
    index = r(len(data_set))
    env = cv2.imread(data_set[index])
    env = cv2.resize(env, (img.shape[1], img.shape[0]))
    bak = (img == 0);
    bak = bak.astype(np.uint8) * 255;
    inv = cv2.bitwise_and(bak, env)
    img = cv2.bitwise_or(inv, img)
    return img


def GenCh(f, val):
    """
    生成中文字符
    """
    img = Image.new("RGB", (45, 70), (255, 255, 255))
    draw = ImageDraw.Draw(img)
    draw.text((0, 3), val, (0, 0, 0), font=f)
    img = img.resize((23, 70))
    A = np.array(img)
    return A


def GenCh1(f, val):
    """
    生成英文字符
    """
    img = Image.new("RGB", (23, 70), (255, 255, 255))
    draw = ImageDraw.Draw(img)
    draw.text((0, 2), val, (0, 0, 0), font=f)
    A = np.array(img)
    return A


def AddGauss(img, level):
    """
    添加高斯模糊
    """
    return cv2.blur(img, (level * 2 + 1, level * 2 + 1));


def r(val):
    return int(np.random.random() * val)


def AddNoiseSingleChannel(single):
    """
    添加高斯噪声
    """
    diff = 255 - single.max();
    noise = np.random.normal(0, 1 + r(1), single.shape);
    noise = (noise - noise.min()) / (noise.max() - noise.min())
    noise = diff * noise;
    noise = noise.astype(np.uint8)
    dst = single + noise
    return dst


def addNoise(img, sdev=0.5, avg=10):
    img[:, :, 0] = AddNoiseSingleChannel(img[:, :, 0]);
    img[:, :, 1] = AddNoiseSingleChannel(img[:, :, 1]);
    img[:, :, 2] = AddNoiseSingleChannel(img[:, :, 2]);
    return img


class GenPlate:

    def __init__(self, fontCh, fontEng, NoPlates):
        self.fontC = ImageFont.truetype(fontCh, 43, 0);
        self.fontE = ImageFont.truetype(fontEng, 60, 0);
        self.img = np.array(Image.new("RGB", (148, 70), (255, 255, 255)))
        # self.bg = cv2.resize(cv2.imread("../images/template.bmp"), (138, 70));
        # self.smu = cv2.imread("../images/smu2.jpg");
        self.noplates_path = [];
        for parent, parent_folder, filenames in os.walk(NoPlates):
            for filename in filenames:
                path = parent + "/" + filename;
                self.noplates_path.append(path);

    def draw(self, val):
        offset = 2
        # self.img[0:70, offset + 8:offset + 8 + 23] = GenCh(self.fontC, val[0]);
        # self.img[0:70, offset + 8 + 23 + 6:offset + 8 + 23 + 6 + 23] = GenCh1(self.fontE, val[1]);
        for i in range(5):
            base = offset + i * 23 + i * 6;
            self.img[0:70, base: base + 23] = GenCh1(self.fontE, val[i]);
        return self.img

    def generate(self, text):
        # if len(text) == 9:
        if len(text) == 5:
            fg = self.draw(text);
            fg = cv2.bitwise_not(fg);
            # com = cv2.bitwise_or(fg, self.bg);
            # com = rot(com, r(60) - 30, com.shape, 30);
            # com = rotRandrom(com, 10, (com.shape[1], com.shape[0]));
            com = tfactor(fg)
            # com = random_envirment(com, self.noplates_path);
            # com = AddGauss(com, 1 + r(2));
            # com = AddGauss(com, r(2));
            # com = addNoise(com);
            return com

    def genPlateString(self, pos, val):
        '''
	    生成车牌String,存为图片
        生成车牌list,存为label
        '''
        plateStr = "";
        plateList = []
        box = [0, 0, 0, 0, 0];
        if (pos != -1):
            box[pos] = 1;
        for unit, cpos in zip(box, range(len(box))):
            plateStr += chars[r(10)]
            plateList.append(plateStr)
        b = [plateList[i][-1] for i in range(len(plateList))]
        return plateStr, b

    # 将生成的车牌图片写入文件夹，对应的label写入label.txt
    def genBatch(self, batchSize, pos, charRange, outputPath,label_file_name, size):
        if (not os.path.exists(outputPath)):
            os.mkdir(outputPath)
        outfile = open(label_file_name, 'w')
        for i in range(batchSize):
            if i % 1000 == 0:
                print(str(i) + ":finish!")
            plateStr, plate = G.genPlateString(-1, -1)
            img = G.generate(plateStr);
            img = cv2.resize(img, size);
            file_name = str(i).zfill(5) + ".jpg"
            write_path = outputPath + "/" + file_name
            cv2.imwrite(write_path, img);
            label_line_list = [write_path,''.join(plateStr)]
            outfile.write('\t'.join(label_line_list) + "\n")


# G.genBatch(100,2,range(31,65),"./plate_100",(272,72))

if __name__ == '__main__':
    G = GenPlate("../font/platech.ttf", '../font/platechar.ttf', "../NoPlates")
    print('usage :python simple_digits_generator.py batch_size output_path label_file_path')
    G.genBatch(int(sys.argv[1]), 2, range(31, 65), sys.argv[2],sys.argv[3], (272, 72))